package com.example.toy.collector.beam;

import com.example.toy.collector.RpcServer;
import com.example.toy.core.record.ToyRecord;
import lombok.RequiredArgsConstructor;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.PipelineResult;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.values.PCollection;

import java.util.ArrayList;
import java.util.List;

/**
 * LocalPCollection
 *
 * @author shenbing
 * @date 2021/12/17
 */
@RequiredArgsConstructor
public class LocalPCollection extends ArrayList<ToyRecord> {

  private final PCollection<ToyRecord> pCollection;

  public void collect() {
    PipelineOptions options = pCollection.getPipeline().getOptions();
    BeamCollectionSender sender = new BeamCollectionSender();

    pCollection.apply(ParDo.of(sender));

    Pipeline pipeline = pCollection.getPipeline();
    PipelineResult run = pipeline.run();
    run.waitUntilFinish();

    RpcServer rpcServer = RpcServer.getInstance();
    List<ToyRecord> toyRecords = rpcServer.pop(sender.getSender().getId());
    addAll(toyRecords);
  }
}
